from utils import print_timing
import os, sys

def print_error(error=None):
    print error
    sys.exit(1)

def get_output_file_name(fileformat, schema, dbname, alleletype):
    filedict = {}
    if alleletype == None:
        filedict["ped"] = "%s.%s.ped"%(dbname, schema)
        filedict["tped"] = "%s.%s.tped"%(dbname, schema)
    else:
        filedict["ped"] = "%s.%s.%s.ped"%(dbname, schema, alleletype)
        filedict["tped"] = "%s.%s.%s.tped"%(dbname, schema, alleletype)
    filedict["map"] = "%s.%s.map"%(dbname, schema)
    filedict["tfam"] = "%s.%s.tfam"%(dbname, schema)
    return filedict[fileformat]

@print_timing
def make_all_output(dbname, alleletype, mapcol, destdir, jobs, shard, test, recreate):
    from make_output import make_output
    from utils import get_conf, file_exists, pg_url
    from dbutils import get_db_type
    conf = get_conf()
    dbuser = conf["dbuser"]
    password = conf["password"]
    portnum = conf["portnum"]
    dbstring_dbname = "%s://%s:%s@localhost:%s/%s"%(pg_url(), dbuser, password, portnum, dbname)
    dbtype = get_db_type(dbstring_dbname)
    if dbtype != 'affy6' and dbtype != 'illumina':
        sys.exit("error: database type can only be 'affy6' or 'illumina'.")
    datasets = conf[dbtype].sections

    if shard and test:
        schemas = [s+'_shard_test' for s in datasets]
    elif shard:
        schemas = [s+'_shard' for s in datasets]
    elif test:
        schemas = [s+'_test' for s in datasets]
    else:
        schemas = datasets

    from utils import create_engine_wrapper
    from dbutils import schema_exists
    db = create_engine_wrapper(dbstring_dbname)

    for dataset, schema in zip(datasets, schemas):
        if not schema_exists(db, schema):
            sys.exit("dataset %s is listed in the config but the corresponding schema %s is not in database %s"%(dataset, schema, dbname))

    dbtype = get_db_type(dbstring_dbname)
    if dbtype not in conf.sections:
        sys.exit("dbtype must be one of %s.\nWas the database %s created with init_db.py?"%(', '.join(conf.sections), dbstring_dbname))

    filelist = []
    for schema in schemas:
        for format in 'map', 'ped', 'tped', 'tfam':
            filelist.append(get_output_file_name(format, schema, dbname, alleletype))

    #alternative versions
    if recreate and [x for x in filelist if file_exists(x, do_print=False)]:
        [os.remove(x) for x in filelist if file_exists(x, do_print=False)]
    elif [x for x in filelist if file_exists(x)]:
        sys.exit("Exiting...")

    if shard:
        formats = ['map', 'ped']
    else:
        formats = ['map', 'ped', 'tped', 'tfam']

    for dataset in datasets:
        for format in formats:
            make_output(format, [dataset], dbname, alleletype, mapcol, destdir, jobs, shard, test, recreate)

if __name__=='__main__':
    from optparse import OptionParser
    usr = os.popen('whoami').read().strip()
    usage = "usage: %prog dbname"
    parser = OptionParser(usage=usage)
    parser.add_option("-a", "--allele-type", action="store", dest="alleletype", help="allele type (for illumina only): foward, top")
    parser.add_option("-c", "--map-col", action="store_true", dest="mapcol", help="add column to map file")
    parser.add_option("-d", "--destdir", action="store", type="string", dest="destdir", help="destination directory")
    parser.add_option("-j", "--jobs", action="store", type="int", dest="jobs", help="number of jobs to run in parallel", default=1)
    parser.add_option("-s", "--sharding", action="store_true", dest="shard", help="use geno table sharding")
    parser.add_option("-t", "--test", action="store_true", dest="test", default=False, help="test cases")
    parser.add_option("-y", "--recreate", action="store_true", dest="recreate", help="delete file and recreate")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        print "incorrect number of arguments"
        parser.print_help()
        sys.exit()
    dbname = args[0]
    print "using database %s"%(dbname)
    make_all_output(dbname, options.alleletype, options.mapcol, options.destdir, options.jobs, options.shard, options.test, options.recreate)
